import os
import shutil
from pathlib import Path

def move_matching_files():
    # 源路径：训练集（包含.raw.h5文件）
    source_dir = Path(r"I:/dataset/lastdata/H5/train1-H5/normal")
    
    # 参考路径：已有的验证集
    reference_dir = Path(r"I:/dataset/nuclear_power/nuclear_H5_process/val")
    
    # 目标路径：新的验证集
    target_dir = Path(r"I:/dataset/lastdata/H5/val/normal")
    
    # 确保目标目录存在
    os.makedirs(target_dir, exist_ok=True)
    
    # 检查目录是否存在
    if not source_dir.exists():
        print(f"错误：源目录不存在：{source_dir}")
        return
    if not reference_dir.exists():
        print(f"错误：参考目录不存在：{reference_dir}")
        return
    
    # 获取参考目录中的所有文件夹名称
    reference_folders = [d.name for d in reference_dir.iterdir() if d.is_dir()]
    print(f"参考目录中的文件夹数量: {len(reference_folders)}")
    
    # 获取源目录中的所有.raw.h5文件
    source_files = [f for f in source_dir.iterdir() if f.is_file() and f.name.endswith('.raw.h5')]
    print(f"源目录中的.raw.h5文件数量: {len(source_files)}")
    
    # 记录移动的文件
    moved_files = []
    
    # 移动匹配的文件
    for file in source_files:
        # 获取文件名（去掉.raw.h5后缀）
        file_name_without_ext = file.name.replace('.raw.h5', '')
        
        # 检查文件名是否与参考文件夹名匹配
        if file_name_without_ext in reference_folders:
            target_path = target_dir / file.name
            print(f"移动文件: {file.name}")
            shutil.move(str(file), str(target_path))
            moved_files.append(file.name)
    
    print(f"总共移动了 {len(moved_files)} 个文件")
    print("移动的文件:", moved_files)

if __name__ == "__main__":
    move_matching_files()